import math
import time
from collections import deque
from typing import Callable
import cvxpy
import numpy as np
from spline_path import calc_spline_course
from qpsolvers import solve_qp
from dora import DoraStatus
from numpy import linalg as LA
from scipy.spatial.transform import Rotation as ROT
from sklearn.metrics import pairwise_distances

#set run_code=1 to run navigation through moving obstacles
#set run_code=2 to run only path following using MPC
#set run_code=3 to run only potential field
run_code = 1
LON_CONTROL_START_ASSIST_THROTTLE =  0.3
BRAKE_MAX = 1.0
THROTTLE_MAX = 0.5

NX = 4  # x = x, y, v, yaw
NU = 2  # a = [accel, steer]
T = 5  # horizon length

import pickle

#load model
with open('./Calibrator.pkl','rb') as model:
    Calibrator = pickle.load(model)

# mpc parameters
R = np.diag([0.01, 0.01])  # input cost matrix
Rd = np.diag([0.01, 1.0])  # input difference cost matrix
Q = np.diag([1.0, 1.0, 0.5, 0.5])  # state cost matrix
Qf = Q  # state final matrix
GOAL_DIS = 1.5  # goal distance
STOP_SPEED = 0.5 / 3.6  # stop speed

# iterative paramter
MAX_ITER = 3  # Max iteration
DU_TH = 0.1  # iteration finish param

TARGET_SPEED = 18.0 / 3.6  # [m/s] target speed
N_IND_SEARCH = 10  # Search index number

DT = 0.2  # [s] time tick

# Vehicle parameters
LENGTH = 4.5  # [m]
WIDTH = 2.0  # [m]
BACKTOWHEEL = 1.0  # [m]
WHEEL_LEN = 0.3  # [m]
WHEEL_WIDTH = 0.2  # [m]
TREAD = 0.7  # [m]
WB = 2.85  # [m]

MAX_STEER = np.deg2rad(45.0)  # maximum steering angle [rad]
MAX_DSTEER = np.deg2rad(30.0)  # maximum steering speed [rad/s]
MAX_SPEED = 55.0 / 3.6  # maximum speed [m/s]
MIN_SPEED = -20.0 / 3.6  # minimum speed [m/s]
MAX_ACCEL = 1.0  # maximum accel [m/ss]

show_animation = True

class State:
    """
    vehicle state class
    """

    def __init__(self, x=0.0, y=0.0, yaw=0.0, v=0.0):
        self.x = x
        self.y = y
        self.yaw = yaw
        self.v = v
        self.predelta = None

def pi_2_pi(angle):
    while(angle > math.pi):
        angle = angle - 2.0 * math.pi

    while(angle < -math.pi):
        angle = angle + 2.0 * math.pi

    return angle


def get_linear_model_matrix(v, phi, delta):

    A = np.zeros((NX, NX))
    A[0, 0] = 1.0
    A[1, 1] = 1.0
    A[2, 2] = 1.0
    A[3, 3] = 1.0
    A[0, 2] = DT * math.cos(phi)
    A[0, 3] = - DT * v * math.sin(phi)
    A[1, 2] = DT * math.sin(phi)
    A[1, 3] = DT * v * math.cos(phi)
    A[3, 2] = DT * math.tan(delta) / WB

    B = np.zeros((NX, NU))
    B[2, 0] = DT
    B[3, 1] = DT * v / (WB * math.cos(delta) ** 2)

    C = np.zeros(NX)
    C[0] = DT * v * math.sin(phi) * phi
    C[1] = - DT * v * math.cos(phi) * phi
    C[3] = - DT * v * delta / (WB * math.cos(delta) ** 2)

    return A, B, C


def update_state(state, a, delta):

    # input check
    if delta >= MAX_STEER:
        delta = MAX_STEER
    elif delta <= -MAX_STEER:
        delta = -MAX_STEER

    state.x = state.x + state.v * math.cos(state.yaw) * DT
    state.y = state.y + state.v * math.sin(state.yaw) * DT
    state.yaw = state.yaw + state.v / WB * math.tan(delta) * DT
    state.v = state.v + a * DT

    if state. v > MAX_SPEED:
        state.v = MAX_SPEED
    elif state. v < MIN_SPEED:
        state.v = MIN_SPEED

    return state


def get_nparray_from_matrix(x):
    return np.array(x).flatten()


def calc_nearest_index(state, cx, cy, cyaw, pind):

    dx = [state.x - icx for icx in cx[pind:(pind + N_IND_SEARCH)]]
    dy = [state.y - icy for icy in cy[pind:(pind + N_IND_SEARCH)]]

    d = [idx ** 2 + idy ** 2 for (idx, idy) in zip(dx, dy)]

    mind = min(d)

    ind = d.index(mind) + pind

    mind = math.sqrt(mind)

    dxl = cx[ind] - state.x
    dyl = cy[ind] - state.y

    angle = pi_2_pi(cyaw[ind] - math.atan2(dyl, dxl))
    if angle < 0:
        mind *= -1

    return ind, mind


def predict_motion(x0, oa, od, xref):
    xbar = xref * 0.0
    for i, _ in enumerate(x0):
        xbar[i, 0] = x0[i]

    state = State(x=x0[0], y=x0[1], yaw=x0[3], v=x0[2])
    for (ai, di, i) in zip(oa, od, range(1, T + 1)):
        state = update_state(state, ai, di)
        xbar[0, i] = state.x
        xbar[1, i] = state.y
        xbar[2, i] = state.v
        xbar[3, i] = state.yaw

    return xbar


def iterative_linear_mpc_control(xref, x0, dref, oa, od):
    """
    MPC contorl with updating operational point iteraitvely
    """

    if oa is None or od is None:
        oa = [0.0] * T
        od = [0.0] * T

    for i in range(MAX_ITER):
        xbar = predict_motion(x0, oa, od, xref)
        poa, pod = oa[:], od[:]
        oa, od, ox, oy, oyaw, ov = linear_mpc_control(xref, xbar, x0, dref)
        du = sum(abs(oa - poa)) + sum(abs(od - pod))  # calc u change value
        if du <= DU_TH:
            break
    else:
        print("Iterative is max iter")

    return oa, od, ox, oy, oyaw, ov


def linear_mpc_control(xref, xbar, x0, dref):
    """
    linear mpc control

    xref: reference point
    xbar: operational point
    x0: initial state
    dref: reference steer angle
    """

    x = cvxpy.Variable((NX, T + 1))
    u = cvxpy.Variable((NU, T))

    cost = 0.0
    constraints = []

    for t in range(T):
        cost += cvxpy.quad_form(u[:, t], R)

        if t != 0:
            cost += cvxpy.quad_form(xref[:, t] - x[:, t], Q)

        A, B, C = get_linear_model_matrix(
            xbar[2, t], xbar[3, t], dref[0, t])
        constraints += [x[:, t + 1] == A @ x[:, t] + B @ u[:, t] + C]

        if t < (T - 1):
            cost += cvxpy.quad_form(u[:, t + 1] - u[:, t], Rd)
            constraints += [cvxpy.abs(u[1, t + 1] - u[1, t]) <=
                            MAX_DSTEER * DT]

    cost += cvxpy.quad_form(xref[:, T] - x[:, T], Qf)

    constraints += [x[:, 0] == x0]
    constraints += [x[2, :] <= MAX_SPEED]
    constraints += [x[2, :] >= MIN_SPEED]
    constraints += [cvxpy.abs(u[0, :]) <= MAX_ACCEL]
    constraints += [cvxpy.abs(u[1, :]) <= MAX_STEER]

    prob = cvxpy.Problem(cvxpy.Minimize(cost), constraints)
    prob.solve(solver=cvxpy.ECOS, verbose=False)

    if prob.status == cvxpy.OPTIMAL or prob.status == cvxpy.OPTIMAL_INACCURATE:
        ox = get_nparray_from_matrix(x.value[0, :])
        oy = get_nparray_from_matrix(x.value[1, :])
        ov = get_nparray_from_matrix(x.value[2, :])
        oyaw = get_nparray_from_matrix(x.value[3, :])
        oa = get_nparray_from_matrix(u.value[0, :])
        odelta = get_nparray_from_matrix(u.value[1, :])

    else:
        print("Error: Cannot solve mpc..")
        oa, odelta, ox, oy, oyaw, ov = None, None, None, None, None, None

    return oa, odelta, ox, oy, oyaw, ov


def calc_ref_trajectory(state, cx, cy, cyaw, ck, sp, dl, pind):
    xref = np.zeros((NX, T + 1))
    dref = np.zeros((1, T + 1))
    ncourse = len(cx)

    ind, _ = calc_nearest_index(state, cx, cy, cyaw, pind)

    if pind >= ind:
        ind = pind

    xref[0, 0] = cx[ind]
    xref[1, 0] = cy[ind]
    xref[2, 0] = sp[ind]
    xref[3, 0] = cyaw[ind]
    dref[0, 0] = 0.0  # steer operational point should be 0

    travel = 0.0

    for i in range(T + 1):
        travel += abs(state.v) * DT
        dind = int(np.round(travel / dl))

        if (ind + dind) < ncourse:
            xref[0, i] = cx[ind + dind]
            xref[1, i] = cy[ind + dind]
            xref[2, i] = sp[ind + dind]
            xref[3, i] = cyaw[ind + dind]
            dref[0, i] = 0.0
        else:
            xref[0, i] = cx[ncourse - 1]
            xref[1, i] = cy[ncourse - 1]
            xref[2, i] = sp[ncourse - 1]
            xref[3, i] = cyaw[ncourse - 1]
            dref[0, i] = 0.0

    return xref, ind, dref

def smooth_yaw(yaw):

    for i in range(len(yaw) - 1):
        dyaw = yaw[i + 1] - yaw[i]

        while dyaw >= math.pi / 2.0:
            yaw[i + 1] -= math.pi * 2.0
            dyaw = yaw[i + 1] - yaw[i]

        while dyaw <= -math.pi / 2.0:
            yaw[i + 1] += math.pi * 2.0
            dyaw = yaw[i + 1] - yaw[i]

    return yaw

def do_simulation(cx, cy, cyaw, ck, sp, dl, initial_state):
    """
    Simulation

    cx: course x position list
    cy: course y position list
    cy: course yaw position list
    ck: course curvature list
    sp: speed profile
    dl: course tick [m]

    """

    # goal = [cx[-1], cy[-1]]

    state = initial_state

    # initial yaw compensation
    if state.yaw - cyaw[0] >= math.pi:
        state.yaw -= math.pi * 2.0
    elif state.yaw - cyaw[0] <= -math.pi:
        state.yaw += math.pi * 2.0

    time = 0.0
    x = [state.x]
    y = [state.y]
    yaw = [state.yaw]
    v = [state.v]
    t = [0.0]
    d = [0.0]
    a = [0.0]
    target_ind, _ = calc_nearest_index(state, cx, cy, cyaw, 0)

    odelta, oa = None, None

    cyaw = smooth_yaw(cyaw)
    xref, target_ind, dref = calc_ref_trajectory(
            state, cx, cy, cyaw, ck, sp, dl, target_ind)
        
#         print("===============xxxxxxxxxxx", state.v, type(state.v))

    x0 = [state.x, state.y, state.v, state.yaw]  # current state

    oa, odelta, ox, oy, oyaw, ov = iterative_linear_mpc_control(
            xref, x0, dref, oa, odelta)
        
#         print("==========oxxxxx===xxxxxxxxx", ox, type(ox))
#         print("xxxxxxxxxxxx", oyaw, type(oyaw))

    if odelta is not None:
        di, ai = odelta[0], oa[0]

    return oa,odelta#t, x, y, yaw, v, d, a

def calc_speed_profile(cx, cy, cyaw, target_speed):

    speed_profile = [target_speed] * len(cx)
    direction = 1.0  # forward

    # Set stop point
    for i in range(len(cx) - 1):
        dx = cx[i + 1] - cx[i]
        dy = cy[i + 1] - cy[i]

        move_direction = math.atan2(dy, dx)

        if dx != 0.0 and dy != 0.0:
            dangle = abs(pi_2_pi(move_direction - cyaw[i]))
            if dangle >= math.pi / 4.0:
                direction = -1.0
            else:
                direction = 1.0

        if direction != 1.0:
            speed_profile[i] = - target_speed
        else:
            speed_profile[i] = target_speed

    speed_profile[-1] = 0.0

    return speed_profile

def compute_throttle_and_brake(current_speed: float, target_speed: float, oa: float):
    """Computes the throttle/brake required to reach the target speed.

    It uses the longitudinal controller to derive the required information.

    Args:

        pid: The mpc controller.

        current_speed (:obj:`float`): The current speed of the ego vehicle

            (in m/s).

        target_speed (:obj:`float`): The target speed to reach (in m/s).

    Returns:

        Throttle and brake dora_input["data"]s.

    """
    if current_speed < 0:
        print("Current speed is negative: {}".format(current_speed))
        non_negative_speed = 0
    else:
        non_negative_speed = current_speed
    acceleration = oa
    model_input = np.array([[non_negative_speed,acceleration]])
    control = Calibrator.predict(model_input)+LON_CONTROL_START_ASSIST_THROTTLE/2

        if control >= 0.0:
            throttle = min(control, LON_CONTROL_MAX_THROTTLE)
            brake = 0.0
        else:
            throttle = 0.0
            brake = min(abs(control), LON_CONTROL_MAX_BRAKE)
    # Keep the brake pressed when stopped or when sliding back on a hill.

    if (current_speed < 1 and target_speed == 0) or current_speed < -0.3:
        brake = 1.0
    return throttle, brake


class Operator:
    """
    Compute the throttle, target angle and brake given a `position`, a `speed` and a `waypoints`.
    """

    def __init__(self):
        self.waypoints = []
        self.rx = []
        self.ry = []
        self.ryaws = [] 
        self.rks = []
        self.flag = False
        self.target_speeds = []
        self.metadata = {}
        self.position = []
        self.speed = []
        self.previous_position = []
        self.current_speed = []
        self.previous_time = time.time()

    def on_event(
        self,
        dora_event: dict,
        send_output: Callable[[str, bytes], None],
    ) -> DoraStatus:
        if dora_event["type"] == "INPUT":
            return self.on_input(dora_event, send_output)
        return DoraStatus.CONTINUE

    def on_input(
        self,
        dora_input: dict,
        send_output: Callable[[str, bytes], None],
    ):
        """Handle input.
        Args:
            dora_input["id"](str): Id of the input declared in the yaml configuration
            dora_input["data"] (bytes): Bytes message of the input
            send_output (Callable[[str, bytes]]): Function enabling sending output back to dora.
        """

        if "position" == dora_input["id"]:
            self.position = np.frombuffer(dora_input["data"], np.float32)
            return DoraStatus.CONTINUE

        elif dora_input["id"] == "speed":
            self.speed = np.frombuffer(dora_input["data"], np.float32)
            return DoraStatus.CONTINUE

        elif "waypoints" == dora_input["id"]:
            waypoints = np.frombuffer(dora_input["data"], np.float32)
            waypoints = waypoints.reshape((-1, 3))
            # print("==============xxxxxxx", waypoints)
            # self.ryaws =  waypoints[:, 3] 
            # self.rks = waypoints[:, 4]
            self.target_speeds = waypoints[:, 2]
            self.waypoints = waypoints

            if len(self.waypoints) >= 5:
                ref_xs = self.waypoints[:, 0]
                ref_ys = self.waypoints[:, 1]
                rxs, rys, ryaws, rks, ss = self._ref_path_smooth(ref_xs=ref_xs, ref_ys=ref_ys)
                self.rx = [item.item() for item in rxs]
                self.ry = [item.item() for item in rys]
                self.waypoints = np.column_stack((rxs, rys))
                # self.ryaws = ryaws
                # self.rks = rks
                self.waypoints = self.waypoints
                self.ryaws = ryaws
                self.rks = rks
                self.rx = self.rx
                self.ry = self.ry
                # 通过插值的方式将速度长度和waypoints长度对齐
                target_length = len(self.waypoints)
                original_length = len(self.target_speeds)
                original_x = np.arange(original_length)
                target_x = np.linspace(0, original_length - 1, target_length) # 目标横坐标，线性插值到目标长度
                self.target_speeds = np.interp(target_x, original_x, self.target_speeds)
            else:
                print("======================xxxxxxxxxxxxxxxxxxxx", self.waypoints)
                target_length = len(self.waypoints)
                ref_xs = self.waypoints[:, 0]
                ref_ys = self.waypoints[:, 1]
                rks, ryaws = self._calculate_curvatures(self.waypoints)
                self.rx = [item.item() for item in ref_xs]
                self.ry = [item.item() for item in ref_ys]
                self.waypoints = self.waypoints
                self.ryaws = ryaws
                self.rks = rks
                self.rx = self.rx
                self.ry = self.ry
                print("===========xxxxxxxxxxxxxxxxxxxx", self.rks)
                print("===========xxxxxxxxxxxxxxxxxxxx", self.ryaws)
                # 通过插值的方式将速度长度和waypoints长度对齐
                target_length = len(self.waypoints)
                original_length = len(self.target_speeds)
                original_x = np.arange(original_length)
                target_x = np.linspace(0, original_length - 1, target_length) # 目标横坐标，线性插值到目标长度
                self.target_speeds = np.interp(target_x, original_x, self.target_speeds)


            waypoints = self.waypoints
            self.metadata = dora_input["metadata"]

        if len(self.position) == 0 or len(self.speed) == 0:
            return DoraStatus.CONTINUE

        if len(self.waypoints) == 0:
            send_output(
                "control",
                np.array([0, 0, 1], np.float16).tobytes(),
                self.metadata,
            )
            return DoraStatus.CONTINUE

        dl = 1.0
        [x, y, _, rx, ry, rz, rw] = self.position
        [_, _, yaw] = ROT.from_quat([rx, ry, rz, rw]).as_euler(
            "xyz", degrees=False
        )
        
        sp = calc_speed_profile(self.rx, self.ry, self.ryaws, self.target_speeds[0]) #TODO:修改TARGET_SPEED
        initial_state = State(x, y, yaw, self.speed[0])
        a,delta = do_simulation(self.rx, self.ry, self.ryaws, self.rks, sp, dl, initial_state)

        if len(self.waypoints) == 0:
            target_angle = 0
            target_speed = 0
        else:
            target_speed = self.target_speeds[0]
            target_angle = delta[0]

        throttle, brake = compute_throttle_and_brake(self.speed[0], target_speed, a[0])

        send_output(
            "control",
            np.array([throttle, target_angle, brake], np.float16).tobytes(),
            self.metadata,
        )
        return DoraStatus.CONTINUE
    
    def _ref_path_smooth(self, ref_xs, ref_ys, max_offset_x=0.5, max_offset_y=0.5, w_ref=1, w_smooth=1.0, w_len=0):
        num_points = len(ref_xs)
        #

        w_ref = w_ref  # weights for close to the ref points

        w_len = w_len  # weights for the points distance

        w_smooth = w_smooth  # smooth of the points


        # init cost matrix

        A1 = np.zeros((num_points * 2 - 4, num_points * 2))
        A2 = np.zeros((num_points * 2 - 2, num_points * 2))
        A3 = np.eye(num_points * 2)
        h = np.zeros(num_points * 2)
        ub = np.zeros(num_points * 2)
        lb = np.zeros(num_points * 2)

        # build up A1 for cost of smooth

        for i in range(2 * num_points - 5):
            A1[i][i] = 1
            A1[i][i + 2] = -2
            A1[i][i + 4] = 1

            A1[i + 1][i + 1] = 1
            A1[i + 1][i + 3] = -2
            A1[i + 1][i + 5] = 1

        # build A2 for distance to the neighbouring points

        for i in range(2 * num_points - 3):
            A2[i][i] = 1
            A2[i][i + 2] = -1

            A2[i + 1][i + 1] = 1
            A2[i + 1][i + 3] = -1

        # A3 is just identity matix


        # A4 and b is for some anchoring points: start two and end two points

        A4 = np.zeros_like(A3)
        A4[0, 0] = 1
        A4[1, 1] = 1
        A4[2, 2] = 1
        A4[3, 3] = 1

        A4[-3, -3] = 1
        A4[-4, -4] = 1
        A4[-1, -1] = 1
        A4[-2, -2] = 1
        b = np.zeros(num_points * 2)
        b[0] = ref_xs[0]
        b[1] = ref_ys[0]
        b[2] = ref_xs[1]
        b[3] = ref_ys[1]
        b[-4] = ref_xs[-2]
        b[-3] = ref_ys[-2]
        b[-2] = ref_xs[-1]
        b[-1] = ref_ys[-1]

        # build h

        h[0 : 2 * num_points : 2] = ref_xs[:]
        h[1 : 2 * num_points : 2] = ref_ys[:]

        # x

        lb[0 : 2 * num_points : 2] = ref_xs[:] - max_offset_x
        ub[0 : 2 * num_points : 2] = ref_xs[:] + max_offset_x
        # y

        lb[1 : 2 * num_points : 2] = ref_ys[:] - max_offset_y
        ub[1 : 2 * num_points : 2] = ref_ys[:] + max_offset_y

        # build quadratic programming

        H = (A1.T @ A1 * w_smooth + A2.T @ A2 * w_len + A3 * w_ref) * 2.0
        f = -w_ref * h * 2

        #  solve the QuadProg daqp, ecos, osqp, scs

        X = solve_qp(H, f, A=A4, b=b, lb=lb, ub=ub, solver='daqp')

        smooth_xs = X[0 : num_points * 2 : 2]
        smooth_ys = X[1 : num_points * 2 + 1 : 2]

        # format with cubic spline

        rxs, rys, ryaws, rks, ss = calc_spline_course(smooth_xs, smooth_ys)

        # print("======================First==================", ryaws)

        # ryaws = self._process_angle(ryaws)

        return rxs, rys, ryaws, rks, ss
    
    def _calculate_curvatures(self, waypoints):
        # 计算路径点的曲率

        num_waypoints = len(waypoints)
        curvatures = np.zeros(num_waypoints)
        yaws = np.zeros(num_waypoints)

        for i in range(1, num_waypoints-1):
            dx1 = waypoints[i][0] - waypoints[i-1][0]
            dy1 = waypoints[i][1] - waypoints[i-1][1]
            dx2 = waypoints[i+1][0] - waypoints[i][0]
            dy2 = waypoints[i+1][1] - waypoints[i][1]

            curvature = abs(dx1 * dy2 - dx2 * dy1) / np.sqrt((dx1**2 + dy1**2) * (dx2**2 + dy2**2 + 1e-8))
            yaws[i] = math.atan2(dy1, dx1)

            #TODO: new yaws
            # yaws[i] = 

            curvatures[i] = curvature

        # yaws = self.

        return curvatures, yaws

    def wrap_angle(self, angle):
        """
        Wrap angle between -pi to pi
        """
        wrap_angle = (angle + math.pi) % (2 * math.pi) - math.pi

        return wrap_angle

    def convert_angle_to_monotonic(self, raw_angles):
        """
        Convert euler angle including +/- 2pi to 0 jump to continous series data
        """
        if len(raw_angles) <= 1:
            return np.copy(raw_angles)

        angle_monotonic = np.zeros(len(raw_angles))

        angle_monotonic[0] = raw_angles[0]
        for i in range(1, len(raw_angles)):
            difference = raw_angles[i] - raw_angles[i - 1]
            angle_monotonic[i] = angle_monotonic[i - 1] + self.wrap_angle(difference)

        return angle_monotonic
        

    def _process_angle(self, raw_angles):
        """
        Adjust angels to start from [-pi, pi] and change monotonically
        """
        adjusted_angles = np.zeros_like(raw_angles)
        for i in range(len(adjusted_angles)):
            adjusted_angles[i] = self.wrap_angle(raw_angles[i])

        adjusted_angles = self.convert_angle_to_monotonic(adjusted_angles)

        return adjusted_angles
